set.seed(1782318297)

library(tidyverse)

# Results

## baseline confounders
pre_cov <- c("age", "female", "education", "income", "race", "urban",
             "country")

##Mediator and outcome equations

m_form <- sat ~ age + female + education + income + race + urban + country + vic + aoj11 + wave

y_form <- sup ~ sat + age + female + education + income + race + urban + country + vic + 
  aoj11*sat + wave


## Models for the post-treatment confounder 

m1 <- lm(aoj11 ~ age + female + education + income + race + urban + country + vic, weights = weights,
         data = nn_data)

m2 <- lm(aoj11 ~ age + female + education + income + race + urban + country + vic, weights = weights,
         data = full_data)

m3 <- lm(aoj11 ~ age + female + education + income + race + urban + country + vic, weights = weights,
         data = cem_data)

## RWR estimation

install.packages("remotes")
remotes::install_github("xiangzhou09/rwrmed")

library(rwrmed)

fit_nn <- rwrmed(treatment = "vic", pre_cov = pre_cov, zmodels = list(m1),
                 y_form = y_form, m_form = m_form, weights = weights, data = nn_data)

fit_full <- rwrmed(treatment = "vic", pre_cov = pre_cov, zmodels = list(m2),
                   y_form = y_form, m_form = m_form, weights = weights, data = full_data)

fit_cem <- rwrmed(treatment = "vic", pre_cov = pre_cov, zmodels = list(m3),
                  y_form = y_form, m_form = m_form, weights = weights, data = cem_data)

##Effect decomposition


###Nearest Neighbor


out_nn <- decomp(fit_nn, rep = 500)

print(out_nn, digits = 2)

twocomp_nn <- out_nn[["twocomp"]]

twocomp_nn <- as.data.frame(twocomp_nn)

twocomp_nn <- twocomp_nn %>% 
  rownames_to_column(var = "effect_type") %>% 
  mutate(match = "Nearest Neighbor")

###Full Matching

out_full <- decomp(fit_full, rep = 500)
print(out_full, digits = 2)

twocomp_full <- out_full[["twocomp"]]

twocomp_full <- as.data.frame(twocomp_full)

twocomp_full <- twocomp_full %>% 
  rownames_to_column(var = "effect_type") %>% 
  mutate(match = "Full Matching")


###CEM

out_cem <- decomp(fit_cem, rep = 500)
print(out_cem, digits = 2)

twocomp_cem <- out_cem[["twocomp"]]

twocomp_cem <- as.data.frame(twocomp_cem)

twocomp_cem <- twocomp_cem %>% 
  rownames_to_column(var = "effect_type") %>% 
  mutate(match = "CEM")

###Figure 3

twocomp <- rbind(twocomp_nn, twocomp_full, twocomp_cem)

library(showtext)

font_add("Cambria", "cambria.ttc")

font_families()

showtext_auto()

twocomp_plot <- twocomp %>% 
  ggplot(aes(x = Estimate, y = effect_type)) +
  geom_point() +
  geom_pointrange(aes(xmin = `2.5% Perc`, xmax = `97.5% Perc`)) +
  geom_vline(xintercept = 0, linetype = "dashed", color = "black") +
  labs(y = "", x = "") +
  theme_minimal() +
  theme(text=element_text(size=11,  family= "Cambria")) +
  facet_grid(~match)
